"""
    jupylet/shadertoy.py
    
    Copyright (c) 2020, Nir Aides - nir@winpdb.org

    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions are met:

    1. Redistributions of source code must retain the above copyright notice, this
       list of conditions and the following disclaimer.
    2. Redistributions in binary form must reproduce the above copyright notice,
       this list of conditions and the following disclaimer in the documentation
       and/or other materials provided with the distribution.

    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
    ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
    WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
    DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
    ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
    (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
    LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
    ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
    SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""


import functools
import webcolors
import datetime
import moderngl
import pathlib
import weakref
import math
import time
import glm
import re

import scipy.signal
import PIL.Image

import moderngl_window as mglw
import numpy as np

from moderngl_window.meta import ProgramDescription
from moderngl_window.opengl import program

from .audio.device import get_output_as_array
from .audio import FPS

from .resource import load_texture, pil_from_texture, find_path, get_context
from .utils import glm_dumps, glm_loads
from .color import c2v
from .state import State
from .node import Node, aa2q, q2aa
from .env import get_window_size
from .lru import SPRITE_TEXTURE_UNIT


@functools.lru_cache()
def get_indices(size):
    
    i0 = np.arange(size)
    i1 = np.ones((size, size))
    i2 = (i1 * i0 + i0[:,None]).astype('long')
    
    return i2


def get_correlation(a0, size=300, step=2, prev=[]):
    
    size = min(size, len(a0))
    
    if not prev:
        prev.append(a0)
        return 0

    p0 = prev[0]

    s0 = size // 2 // step
    c0 = a0[get_indices(s0) * step] @ p0[::step][:s0]
    
    ix = c0.argmax() * step
    a1 = a0[ix:]
    
    prev[0] = a1

    return ix


def get_shadertoy_audio(
    amp=1., 
    length=512, 
    buffer=500, 
    data=None, 
    channel_time=None,
    correlate=True,
    ):
    
    if data is not None:
        a0 = data
        ct = channel_time

    else:
        l0 = length + buffer
        a0, ct = get_output_as_array(-l0, l0)[:2]
        
        if a0 is None:
            a0 = np.zeros((length, 2))

    if correlate:
        ix = get_correlation(a0.mean(-1), buffer)
        a0 = a0[ix:][:length]
            
    if channel_time is not None:
        ct = channel_time

    ft = np.fft.rfft(a0, axis=0)
    sa = np.square(np.abs(ft)) + 1e-6
    ps = 10 * np.log10(sa)
    rs = scipy.signal.resample(ps, len(a0))
    ns = (rs + 50) / 100
    
    return np.stack((ns * 256, a0 * amp * 128 + 128)).clip(0, 255), ct


def load_shadertoy_program(source):

    if '\n' not in source:
        path = find_path(source)
        source = path.open().read()

    path = find_path('shaders/shadertoy-wrapper.glsl')
    
    single = path.open().read()
    single = re.sub(r'void mainImage.* {}', source, single)

    pd = ProgramDescription(path='shadertoy.glsl')
    sd = program.ProgramShaders.from_single(pd, single)
    pg = sd.create()

    return pg


class Shadertoy(Node):

    """一个Shadertoy画布(canvas)。
    
    Args:
        width (float): shadertoy画布的宽度。
        height (float): shadertoy画布的高度。
        x (float): shadertoy画布位置的x坐标量。
        y (float): shadertoy画布位置的y坐标量。
        angle (float): 以度为单位顺时针旋转shadertoy画布的量。
        anchor_x (float or str):  'left', 'center' ， 'right' 或介于0.0（左）和
            1.0（右）之间的值，指示shadertoy画布内沿其x轴的定位点。
        anchor_y (float or str):  'bottom', 'center' ， 'top' 或介于0.0（底）和
            1.0（顶）之间的值，指示shadertoy画布内沿其y轴的定位点。
        color (str or 3-tuple): 用于着色shadertoy画布图像的颜色。可以是颜色名称、颜色十六进制符号或三元组(3-tuple)。
    """

    def __init__(
        self,
        code,
        width=800,
        height=450,
        x=0, 
        y=0,
        angle=0.0,
        anchor_x='left',
        anchor_y='bottom',
        color='white',
        name=None,
    ):
        """"""

        super().__init__(
            name,
            rotation=aa2q(glm.radians(angle)),
            position=glm.vec3(x, y, 0),
        )

        self.t0 = None
        self.ct = 0
        self.dt = 0

        self.iframe = 0

        w0, h0 = get_window_size()
        
        self.shader = load_shadertoy_program(code)
        self.shader._members['jpl_projection'].write(glm.ortho(
            0, w0, 0, h0, -1, 1
        ))

        self.geometry = mglw.geometry.quad_2d(
            size=(1.0, 1.0), 
            pos=(0.5, 0.5)
        )

        self.channel0 = None
        self.channel1 = None
        self.channel2 = None
        self.channel3 = None
        
        self.channeltime = [0., 0., 0., 0.]

        self.width = width
        self.height = height
        self.components = 4
        
        self.color4 = glm.vec4(1., 1., 1., 1.)

        self.set_anchor(anchor_x, anchor_y)
        self.color = color

        self.tx0 = None
        self.tx1 = None
        self.fbo = None

    def __del__(self):
        self.release()

    def release(self):

        if self.tx0 is not None:

            self.tx0.release()
            self.tx0 = None

            self.tx1.release()
            self.tx1 = None

        if self.fbo is not None:
            self.fbo.release()
            self.fbo = None

    def update(self, shader):
        pass

    def use(self, location):
        self.tx0.use(location=location)

    def render2buffer(self, ct, dt, iframe, width, height):

        if self.iframe == iframe:
            return

        ctx = get_context()
        fb0 = ctx.fbo

        if (self.width, self.height) != (width, height):

            self.release()

        if self.tx0 is None:

            self.width = width
            self.height = height

            self.tx0 = ctx.texture((int(self.width), int(self.height)), self.components, dtype='f4')        
            self.tx1 = ctx.texture((int(self.width), int(self.height)), self.components, dtype='f4')

        self.fbo is not None and self.fbo.release()
        self.fbo = ctx.framebuffer(color_attachments=[self.tx1])
        self.fbo.use()
        self.fbo.clear()

        self.shader._members['jpl_projection'].write(glm.ortho(
            0, self.width, 0, self.height, -1, 1
        ))

        self.render(ct, dt)

        self.tx0, self.tx1 = self.tx1, self.tx0

        fb0.use()

    def draw(self, ct, dt):
        """将shadertoy渲染到画布-这是 shadertoy.render() 的别名"""
        return self.render(ct, dt)
        
    def render(self, ct, dt):
        """将shadertoy渲染到画布。"""
        
        if self.t0 is None:
            self.t0 = ct

        self.ct = ct - self.t0
        self.dt = dt

        self.iframe += 1

        if self._dirty:
            self.update(self.shader)

        for i in [0, 1, 2, 3]:
            ch = getattr(self, 'channel%s' % i, None)
            if isinstance(ch, Shadertoy):
                ch.render2buffer(ct, dt, self.iframe, self.width, self.height)

        self.shader._members['jpl_model'].write(self.matrix)
        self.shader._members['jpl_components'].value = self.components
        self.shader._members['jpl_color'].write(self.color4)

        if 'iResolution' in self.shader._members:
            self.shader._members['iResolution'].write(self.scale)
        
        if 'iTime' in self.shader._members:
            self.shader._members['iTime'].value = self.ct
        
        if 'iTimeDelta' in self.shader._members:
            self.shader._members['iTimeDelta'].value = dt

        if 'iFrame' in self.shader._members:
            self.shader._members['iFrame'].value = self.iframe

        for i in [0, 1, 2, 3]:
            
            ch = getattr(self, 'channel%s' % i, None)
            if ch is None:
                continue

            if 'iChannel%s' % i in self.shader._members:
                self.shader._members['iChannel%s' % i].value = SPRITE_TEXTURE_UNIT + i
                ch.use(location=SPRITE_TEXTURE_UNIT+i)

            if 'iChannelTime[%s]' % i in self.shader._members:
                self.shader._members['iChannelTime[%s]' % i].value = self.channeltime[i]

            if 'iChannelResolution[%s]' % i in self.shader._members:
                self.shader._members['iChannelResolution[%s]' % i].write(
                    glm.vec3(ch.width, ch.height, ch.components)
                )

        if 'iDate' in self.shader._members:
            dt = datetime.datetime.now()
            self.shader._members['iDate'].write(glm.vec4(dt.year, dt.month, dt.day, time.time()))

        if 'iSampleRate' in self.shader._members:
            self.shader._members['iSampleRate'].value = FPS

        self.geometry.render(self.shader)

    @property
    def x(self):
        """float: 锚定位置的x坐标。"""
        return self.position.x
        
    @x.setter
    def x(self, value):
        self.position.x = value
        
    @property
    def y(self):
        """float: 锚定位置的y坐标。"""
        return self.position.y
        
    @y.setter
    def y(self, value):
        self.position.y = value
        
    @property
    def angle(self):
        """float: 围绕锚点的旋转角度（以度为单位）。"""
        angle, axis = q2aa(self.rotation)
        return round(glm.degrees(angle * glm.sign(axis.z)), 4)

    @angle.setter
    def angle(self, angle):
        self.rotation = aa2q(glm.radians(angle))

    def set_anchor(self, ax=None, ay=None):
        """设置shadertoy画布的锚点。

        锚点是shadertoy画布中用于旋转和定位的点。想象一个大头针穿过画布，你用
        这个大头针定位画布并旋转它。大头针穿过画布的点是锚定点。

        分别设置锚点的x轴(定位)和y轴(定位)。

        Args:
            ax (str or float): x定位可以是 'left', 'center', 'right' 或
                介于0.0（左）和1.0（右）之间的浮点值之一。

            ay (str or float): y定位可以是 'bottom', 'center', 'top' 或
                介于0.0（底部）和1.0（顶部）之间的浮动值之一。
        """
        self._ax = ax
        self._ay = ay

        if ax == 'left':
            self.anchor.x = 0
        elif ax == 'center':
            self.anchor.x = 0.5
        elif ax == 'right':
            self.anchor.x = 1.
        elif type(ax) in (int, float):
            self.anchor.x = ax / self.width

        if ay == 'bottom':
            self.anchor.y = 0
        elif ay == 'center':
            self.anchor.y = 0.5
        elif ay == 'top':
            self.anchor.y = 1.
        elif type(ay) in (int, float):
            self.anchor.y = ay / self.width

    @property
    def width(self):
        """float: 以像素为单位的宽度。"""
        return self.scale0.x

    @width.setter
    def width(self, width):
        self.scale0.x = width

    @property
    def height(self):
        """float: 以像素为单位的高度。"""
        return self.scale0.y

    @height.setter
    def height(self, height):
        self.scale0.y = height

    def set_channel(
        self, 
        channel, 
        data,
        channeltime=None,
        ):

        assert channel in [0, 1, 2, 3]

        if isinstance(data, Shadertoy):
            if data is self:
                data = weakref.proxy(data)
            setattr(self, 'channel%s' % channel, data)
            return

        if channeltime:
            self.channeltime[channel] = channeltime

        channel = 'channel%s' % channel

        texture = getattr(self, channel, None)
        if texture is not None:
            texture.release()
            
        texture = load_texture(
            data,
            anisotropy=1., 
            autocrop=False,
            mipmap=False, 
            flip=False, 
        )
        texture.repeat_x = False
        texture.repeat_y = False

        setattr(self, channel, texture)

    def distance_to(self, o=None, pos=None):
        """计算到另一个精灵或坐标的距离。

        Args:
            o (Sprite, optional): 其他要计算距离的精灵(Sprite)对象。
            pos (tuple, optional): 计算距离的的(x, y)坐标。

        Returns:
            float: 以像素为单位的距离。
        """
        x, y = pos or (o.position.x, o.position.y)
        
        dx = x - self.position.x
        dy = y - self.position.y

        return (dx ** 2 + dy ** 2) ** 0.5
    
    def angle_to(self, o=None, pos=None):
        """以度为单位计算与另一个精灵或坐标的顺时针角度。

        Args:
            o (Sprite, optional): 其他要计算角度的精灵(Sprite)对象。
            pos (tuple, optional): 计算角度的的(x, y)坐标。

        Returns:
            float: 以度为单位的角度值。
        """

        qd = {
            (True, True): 0,
            (True, False): 180,
            (False, False): 180,
            (False, True): 360,
        }
        
        x, y = pos or (o.position.x, o.position.y)
        
        dx = x - self.position.x
        dy = y - self.position.y

        a0 = math.atan(dy / (dx or 1e-7)) / math.pi * 180 + qd[(dy >= 0, dx >= 0)]

        return -a0

    @property
    def top(self):
        """float: 获取精灵边界框的顶部坐标。"""
        t, r, b, l = self._trbl()
        return self.position.y + t
        
    @property
    def right(self):
        """float: 获取精灵边界框的右部坐标。"""
        t, r, b, l = self._trbl()
        return self.position.x + r
        
    @property
    def bottom(self):
        """float: 获取精灵边界框的底部坐标。"""
        t, r, b, l = self._trbl()
        return self.position.y + b
        
    @property
    def left(self):
        """float: 获取精灵边界框的左部坐标。"""
        t, r, b, l = self._trbl()
        return self.position.x + l
        
    @property
    def radius(self):
        """float: 获取包含精灵边界框的圆的半径。"""
        t, r, b, l = self._trbl()
        rs = max(t, b) ** 2 + max(r, l) ** 2
        return rs ** .5
        
    def _trbl(self):
        return trbl(
            self.width, 
            self.height, 
            self.anchor.x * self.width, 
            self.anchor.y * self.height, 
            self.angle,
        )

    def wrap_position(self, width, height, margin=50):
        """指定精灵环绕画布的宽度和高度。

        如果精灵从画布的另一侧消失，请使用此方法使其从画布的一侧返回

        Args:
            width (float): 要环绕的画布宽度。
            height (float): 要环绕的画布高度。
            margin (float, optional): 在将精灵包裹到另一侧之前，在画布周围添加额外的边距。
        """
        self.position.x = (self.position.x + margin) % (width + 2 * margin) - margin
        self.position.y = (self.position.y + margin) % (height + 2 * margin) - margin

    def clip_position(self, width, height, margin=0):
        """将精灵坐标剪裁到给定的画布宽度和高度。

        使用此方法可防止精灵离开画布。

        Args:
            width (float): 要剪裁到的画布宽度。
            height (float): 要剪裁到的画布高度。
            margin (float, optional): 在剪切精灵坐标之前，在画布周围添加额外的边距。
        """
        self.position.x = max(-margin, min(margin + width, self.position.x))
        self.position.y = max(-margin, min(margin + height, self.position.y))

    @property
    def opacity(self):
        """float: 获取或设置精灵的不透明度(opacity)。

        将opacity设置为0将使精灵完全透明。
        将opacity设置为1将使精灵完全不透明。
        """
        return self.color4.a

    @opacity.setter
    def opacity(self, opacity):
        self.color4.a = opacity
        
    @property
    def color(self):
        """glm.vec4: 获取或设置精灵的颜色。
        
        精灵颜色将乘以位图图像的颜色值。

        颜色可以通过名称(例如 'white')或十六进制表示法(例如 '#cc4488') 
        或 四元组(4-tuple) 或 glm.vec4 值来指定
        """
        return self.color4

    @color.setter
    def color(self, color):        
        self.color4 = c2v(color, self.color4.a)

